"""深度学习垃圾分类"""
from keras.models import load_model
import json
import numpy as np
import cv2

class gar_classif():
    def __init__(self,weight_path,index_path):
        self.weight_path=weight_path
        self.index_path=index_path
        #加载resnet50
        self.model = load_model(self.weight_path)

    def classify(self,img_path):
        """垃圾分类"""
        model=self.model
        with open("./data/index.json", 'r') as load_f:
            load_dict = json.load(load_f)
        imgori=cv2.imread(img_path,cv2.IMREAD_COLOR)
        height, width = imgori.shape[:2]
        #显示原图
        the_img=cv2.resize(imgori, (int(width/6), int(height/6)))
        #原图预处理
        img = cv2.resize(imgori, (224, 224))
        cv2.imshow("the img", the_img)
        img=img.reshape(1,224,224,3)
        #识别
        pred=model.predict(img)
        #根据索引找到分类结果
        label= np.argmax(pred, axis=1)#axis = 1是取行的最大值的索引，0是列的最大值的索引
        print("the label index:",label)
        print("classify result:",load_dict[str(label[0])])
        cv2.waitKey(0)

if __name__ == "__main__":
    weight_path="./models/garclass.h5"
    index_path="./data/index.json"
    img="./data/val_data/cans.jpg"
    net=gar_classif(weight_path,index_path)
    net.classify(img)
